import pandas as pd
import numpy as np
import datetime

stime = '20200101'
etime = '20201230'
stime1 = datetime.datetime.strptime(stime, '%Y%m%d').strftime('%Y-%m-%d')
etime1 = datetime.datetime.strptime(etime, '%Y%m%d').strftime('%Y-%m-%d')

pool = get_index_stocks('399012.SZ', '20200108')


def get_daily_price(group):
    selector1 = ((group.index.hour == 14) & (group.index.minute == 50))
    selector2 = ((group.index.hour == 9) & (group.index.minute == 40))
    selector3 = ((group.index.hour == 9) & (group.index.minute == 45))
    selector4 = ((group.index.hour == 9) & (group.index.minute == 50))
    selector5 = ((group.index.hour == 9) & (group.index.minute == 55))
    return pd.Series([None if group.loc[selector1].empty else group.loc[selector1].close[0],
                      None if group.loc[selector2].empty else group.loc[selector2].close[0],
                      None if group.loc[selector3].empty else group.loc[selector3].close[0],
                      None if group.loc[selector4].empty else group.loc[selector4].close[0],
                      None if group.loc[selector5].empty else group.loc[selector5].close[0]])


def subtract_first(ndarray):
    return ndarray[-1] - ndarray[0]


# df.dropna(subset=['列名'],inplace=True)
all_flows = None
i = 0
for symbol in pool:
    histories_1d = get_price([symbol], stime, etime,
                             '1d', ['close', 'quote_rate', 'turnover', 'turnover_rate'],
                             True,  # 是否跳过停牌
                             "pre",  # 前复权
                             0,  # 天数
                             is_panel=0)[symbol]
    # pd.shift

    histories_1d['3d_up'] = histories_1d.close * 100 / histories_1d.close.shift(3) - 1
    # histories_1d['1d_up'] = histories_1d.close/histories_1d.close.shift(1)-1
    # histories_1d['3d_mean_turnover'] = histories_1d['turnover_rate'].rolling(3).mean()
    # histories_1d["day"] = histories_1d.index
    histories_1d = histories_1d.rename(columns={'close': 'close_1d', 'quote_rate': '1d_up'})

    histories_5m = get_price([symbol], stime + " 09:30", etime + " 15:00", "5m", ['close'], True, 'pre', 0, is_panel=0)[
        symbol]
    histories_5m['day'] = pd.to_datetime(histories_5m.index.date, format='%Y-%m-%d')
    ### 计算出最终得
    res = histories_5m.groupby("day").apply(get_daily_price)
    res.columns = ['14:55', '09:40', '09:45', '09:50', '09:55']

    res["up_940"] = (res['09:40'] / res['14:55'].shift(1)).apply(lambda x: np.round(x - 1, 3) * 100)
    res["up_945"] = (res['09:45'] / res['14:55'].shift(1)).apply(lambda x: np.round(x - 1, 3) * 100)
    res["up_950"] = (res['09:50'] / res['14:55'].shift(1)).apply(lambda x: np.round(x - 1, 3) * 100)
    res["up_955"] = (res['09:55'] / res['14:55'].shift(1)).apply(lambda x: np.round(x - 1, 3) * 100)
    res["up_940"] = res["up_940"].shift(-1)
    res["up_945"] = res["up_945"].shift(-1)
    res["up_950"] = res["up_950"].shift(-1)
    res["up_955"] = res["up_955"].shift(-1)
    res.drop(['09:40', '09:45', '09:50', '09:55'], axis=1, inplace=True)
    flows = get_money_flow_step([symbol], stime1 + " 09:30", etime1 + " 14:50",
                                "5m",
                                ['act_buy_xl', 'act_buy_l', 'act_buy_m', 'act_sell_xl', 'act_sell_l', 'act_sell_m',
                                 'dde_l'],
                                None, is_panel=0)[symbol]
    flows['day'] = pd.to_datetime(flows.index.date, format='%Y-%m-%d')

    flows['act_buy_xl_10'] = flows['act_buy_xl'].rolling(2).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_buy_xl_20'] = flows['act_buy_xl'].rolling(4).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_buy_xl_40'] = flows['act_buy_xl'].rolling(8).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_buy_xl_80'] = flows['act_buy_xl'].rolling(16).sum().apply(lambda x: np.round(np.log(x + 1), 3))

    flows['act_sell_xl_10'] = flows['act_sell_xl'].rolling(2).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_sell_xl_20'] = flows['act_sell_xl'].rolling(4).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_sell_xl_40'] = flows['act_sell_xl'].rolling(8).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_sell_xl_80'] = flows['act_sell_xl'].rolling(16).sum().apply(lambda x: np.round(np.log(x + 1), 3))

    flows['act_buy_l_10'] = flows['act_buy_l'].rolling(2).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_buy_l_20'] = flows['act_buy_l'].rolling(4).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_buy_l_40'] = flows['act_buy_l'].rolling(8).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_buy_l_80'] = flows['act_buy_l'].rolling(16).sum().apply(lambda x: np.round(np.log(x + 1), 3))

    flows['act_sell_l_10'] = flows['act_sell_l'].rolling(2).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_sell_l_20'] = flows['act_sell_l'].rolling(4).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_sell_l_40'] = flows['act_sell_l'].rolling(8).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_sell_l_80'] = flows['act_sell_l'].rolling(16).sum().apply(lambda x: np.round(np.log(x + 1), 3))

    flows['act_buy_m_10'] = flows['act_buy_m'].rolling(2).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_buy_m_20'] = flows['act_buy_m'].rolling(4).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_buy_m_40'] = flows['act_buy_m'].rolling(8).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_buy_m_80'] = flows['act_buy_m'].rolling(16).sum().apply(lambda x: np.round(np.log(x + 1), 3))

    flows['act_sell_m_10'] = flows['act_sell_m'].rolling(2).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_sell_m_20'] = flows['act_sell_m'].rolling(4).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_sell_m_40'] = flows['act_sell_m'].rolling(8).sum().apply(lambda x: np.round(np.log(x + 1), 3))
    flows['act_sell_m_80'] = flows['act_sell_m'].rolling(16).sum().apply(lambda x: np.round(np.log(x + 1), 3))

    flows['dde_10'] = flows['dde_l'].rolling(2).sum().apply(
        lambda x: np.round(np.log(x + 1), 3) if x > 0 else -1 * np.round(np.log(1 - x), 3))
    flows['dde_20'] = flows['dde_l'].rolling(4).sum().apply(
        lambda x: np.round(np.log(x + 1), 3) if x > 0 else -1 * np.round(np.log(1 - x), 3))
    flows['dde_40'] = flows['dde_l'].rolling(8).sum().apply(
        lambda x: np.round(np.log(x + 1), 3) if x > 0 else -1 * np.round(np.log(1 - x), 3))
    flows['dde_80'] = flows['dde_l'].rolling(16).sum().apply(
        lambda x: np.round(np.log(x + 1), 3) if x > 0 else -1 * np.round(np.log(1 - x), 3))

    flows['up_10m'] = (histories_5m.close / histories_5m.close.shift(3)).apply(lambda x: np.round(x - 1, 3) * 100)
    flows['up_20m'] = (histories_5m.close / histories_5m.close.shift(5)).apply(lambda x: np.round(x - 1, 3) * 100)
    flows['up_40m'] = (histories_5m.close / histories_5m.close.shift(9)).apply(lambda x: np.round(x - 1, 3) * 100)
    flows['up_80m'] = (histories_5m.close / histories_5m.close.shift(17)).apply(lambda x: np.round(x - 1, 3) * 100)

    ### 后期进行标准化 dde volume
    flows = flows.loc[((flows.index.hour == 14) & (flows.index.minute == 50))]
    flows.index = flows.day

    flows.drop(['day', 'act_buy_xl', 'act_buy_l', 'act_buy_m', 'act_sell_xl', 'act_sell_l', 'act_sell_m'], axis=1,
               inplace=True)
    flows = flows.join(histories_1d)
    flows = flows.join(res)
    flows['symbol'] = symbol
    flows['date'] = flows.index
    # flows.index == range(flows.s)
    if all_flows is not None:
        all_flows = pd.concat([all_flows, flows], ignore_index=True)
        # all_flows.append(flows)
    else:
        all_flows = flows
    i = i + 1
    print("-----------------", i, symbol)
#     if i > 5:
#         break
all_flows.to_csv('stock_data/all.csv')